import akshare as ak
from clickhouse_driver import Client
from datetime import datetime
import pytz
import pandas as pd

CH = Client(host='clickhouse', port=9000)
CH.execute('CREATE DATABASE IF NOT EXISTS qbot')

CH = Client(host='clickhouse', port=9000, database='qbot')

def create_table():
    CH.execute('''
        CREATE TABLE IF NOT EXISTS tick_1m (
            code String,
            dt DateTime,
            open Float32,
            high Float32,
            low Float32,
            close Float32,
            vol UInt64
        ) ENGINE = MergeTree()
        PARTITION BY toYYYYMMDD(dt)
        ORDER BY (code, dt)
    ''')
    
create_table()   # 自动建表


def crawl_1m(code: str = "sz000001"):          # 默认用平安银行（深圳）
    try:
        # 自动补全前缀
        symbol = code if code.startswith(('sz', 'sh')) else f'sz{code}'
        df = ak.stock_zh_a_minute(symbol=symbol, period='1', adjust='')
        if df is None or df.empty:
            raise ValueError("akshare return empty")
        df = df[['day', 'open', 'high', 'low', 'close', 'volume']]
    except Exception as e:
        print(f"[WARN] akshare failed: {e}  → use fake tick")
        now = datetime.now().replace(second=0, microsecond=0)
        base = round(random.uniform(14, 16), 2)   # 平安银行价格区间
        df = pd.DataFrame([{
            'day': now.strftime('%Y-%m-%d %H:%M:00'),
            'open': base,
            'high': base + random.uniform(0, 0.2),
            'low': base - random.uniform(0, 0.2),
            'close': base + random.uniform(-0.1, 0.1),
            'volume': random.randint(1000, 50000)
        }])

    df.columns = ['dt', 'open', 'high', 'low', 'close', 'vol']
    df['code'] = code
    df['dt'] = pd.to_datetime(df['dt']).dt.tz_localize('Asia/Shanghai')
    df = df.astype({
        'open': 'float32',
        'high': 'float32',
        'low': 'float32',
        'close': 'float32',
        'vol': 'int64'
    })
    rows = [
        (
            str(r['code']),
            r['dt'].to_pydatetime(),
            float(r['open']),
            float(r['high']),
            float(r['low']),
            float(r['close']),
            int(r['vol'])
        )
        for r in df.to_dict('records')
    ]
    
    CH.execute('INSERT INTO tick_1m (code, dt, open, high, low, close, vol) VALUES', rows)
    print(f"[INFO] Insert {len(rows)} rows at {datetime.now()}")

STOCK_POOL = ['sz000001', 'sz000002', 'sh600519']   # 平安+万科+茅台

def crawl_all_1m():
    for code in STOCK_POOL:
        crawl_1m(code)      # 复用已有函数
        
if __name__ == '__main__':
    create_table()
    crawl_1m()
